# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

import logging
import json
from typing import List, Union

from tornado import gen

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase


class GdaxSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               product_id: Union[List[str], str],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._product_id = to_list(product_id)
    self._register_url('wss://ws-feed.pro.coinbase.com')

    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_level2 = topic_map.search_topic_id('gdax_ws_level2_%s%s' % (group_id, worker_id))[0]
    publisher_level2 = queue.get_writer(topic_id=topic_level2[1])

    topic_full = topic_map.search_topic_id('gdax_ws_full_%s%s' % (group_id, worker_id))[0]
    publisher_full = queue.get_writer(topic_id=topic_full[1])

    self._publisher = {
        'snapshot': publisher_level2,
        'l2update': publisher_level2,
        'subscriptions': publisher_full,  # Override level2 subscription.
        'done': publisher_full,
        'received': publisher_full,
        'open': publisher_full,
        'match': publisher_full,
        'margin_profile_update': publisher_full,
        'change': publisher_full,
        'activate': publisher_full,
    }

    self._enable_topic_event_publisher(queue,
                                       'gdax_ws_topic_event', [topic_level2[0], topic_full[0]])

    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  @gen.coroutine
  def on_open(self):
    yield self._ws.write_message(
        json.dumps({
            'type': 'subscribe', 'product_ids': self._product_id, 'channels': ['level2', 'full']
        }))
    self._publish_topic_event('START')

  def on_message(self, timestamp: int, message: dict):
    if self.stopped:
      return
    if 'type' not in message:
      return
    msg_type = message['type']
    try:
      self._publisher[msg_type].write(timestamp, message)
    except KeyError:
      self._logger.error('Unknown type: %s' % msg_type)
